package com.xuanzheng.company.dao.impl;

import com.xuanzheng.company.dao.AIChatDao;
import com.xuanzheng.company.entity.AIChatRecord;
import com.xuanzheng.company.entity.AIUsageLimit;
import com.xuanzheng.company.util.DBUtil;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;

/**
 * AI聊天数据访问实现类
 */
public class AIChatDaoImpl implements AIChatDao {
    
    @Override
    public boolean saveChatRecord(AIChatRecord record) {
        String sql = "INSERT INTO ai_chat_record (id, user_id, user_name, model_id, question, " +
                    "answer, status, error_message, tokens_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)";
        
        try (Connection conn = DBUtil.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            
            pstmt.setString(1, record.getId());
            pstmt.setString(2, record.getUserId());
            pstmt.setString(3, record.getUserName());
            pstmt.setString(4, record.getModelId());
            pstmt.setString(5, record.getQuestion());
            pstmt.setString(6, record.getAnswer());
            pstmt.setInt(7, record.getStatus());
            pstmt.setString(8, record.getErrorMessage());
            pstmt.setInt(9, record.getTokensUsed());
            
            return pstmt.executeUpdate() > 0;
        } catch (SQLException e) {
            e.printStackTrace();
        }
        
        return false;
    }
    
    @Override
    public boolean updateChatRecord(AIChatRecord record) {
        String sql = "UPDATE ai_chat_record SET answer = ?, status = ?, error_message = ?, tokens_used = ? " +
                    "WHERE id = ?";
        
        try (Connection conn = DBUtil.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            
            pstmt.setString(1, record.getAnswer());
            pstmt.setInt(2, record.getStatus());
            pstmt.setString(3, record.getErrorMessage());
            pstmt.setInt(4, record.getTokensUsed());
            pstmt.setString(5, record.getId());
            
            return pstmt.executeUpdate() > 0;
        } catch (SQLException e) {
            e.printStackTrace();
        }
        
        return false;
    }
    
    @Override
    public List<AIChatRecord> findChatRecordsByUserId(String userId, int limit) {
        List<AIChatRecord> records = new ArrayList<>();
        String sql = "SELECT * FROM ai_chat_record WHERE user_id = ? " +
                    "ORDER BY create_time DESC LIMIT ?";
        
        try (Connection conn = DBUtil.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            
            pstmt.setString(1, userId);
            pstmt.setInt(2, limit);
            ResultSet rs = pstmt.executeQuery();
            
            while (rs.next()) {
                records.add(mapResultSetToChatRecord(rs));
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
        
        return records;
    }
    
    @Override
    public AIUsageLimit findUsageLimitByUserId(String userId) {
        String sql = "SELECT * FROM ai_usage_limit WHERE user_id = ?";
        
        try (Connection conn = DBUtil.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            
            pstmt.setString(1, userId);
            ResultSet rs = pstmt.executeQuery();
            
            if (rs.next()) {
                return mapResultSetToUsageLimit(rs);
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
        
        return null;
    }
    
    @Override
    public boolean createUsageLimit(AIUsageLimit limit) {
        String sql = "INSERT INTO ai_usage_limit (id, user_id, daily_limit, used_today, last_reset_date) " +
                    "VALUES (?, ?, ?, ?, ?)";
        
        try (Connection conn = DBUtil.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            
            pstmt.setString(1, limit.getId());
            pstmt.setString(2, limit.getUserId());
            pstmt.setInt(3, limit.getDailyLimit());
            pstmt.setInt(4, limit.getUsedToday());
            pstmt.setDate(5, limit.getLastResetDate());
            
            return pstmt.executeUpdate() > 0;
        } catch (SQLException e) {
            e.printStackTrace();
        }
        
        return false;
    }
    
    @Override
    public boolean updateUsageLimit(AIUsageLimit limit) {
        String sql = "UPDATE ai_usage_limit SET daily_limit = ?, used_today = ?, " +
                    "last_reset_date = ?, update_time = CURRENT_TIMESTAMP WHERE user_id = ?";
        
        try (Connection conn = DBUtil.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            
            pstmt.setInt(1, limit.getDailyLimit());
            pstmt.setInt(2, limit.getUsedToday());
            pstmt.setDate(3, limit.getLastResetDate());
            pstmt.setString(4, limit.getUserId());
            
            return pstmt.executeUpdate() > 0;
        } catch (SQLException e) {
            e.printStackTrace();
        }
        
        return false;
    }
    
    @Override
    public boolean incrementUsageCount(String userId) {
        String sql = "UPDATE ai_usage_limit SET used_today = used_today + 1, " +
                    "update_time = CURRENT_TIMESTAMP WHERE user_id = ?";
        
        try (Connection conn = DBUtil.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            
            pstmt.setString(1, userId);
            return pstmt.executeUpdate() > 0;
        } catch (SQLException e) {
            e.printStackTrace();
        }
        
        return false;
    }
    
    @Override
    public boolean resetDailyUsage() {
        String sql = "UPDATE ai_usage_limit SET used_today = 0, last_reset_date = CURDATE(), " +
                    "update_time = CURRENT_TIMESTAMP WHERE last_reset_date < CURDATE()";
        
        try (Connection conn = DBUtil.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            
            return pstmt.executeUpdate() >= 0; // 可能为0，表示没有需要重置的记录
        } catch (SQLException e) {
            e.printStackTrace();
        }
        
        return false;
    }
    
    /**
     * 将ResultSet映射为AIChatRecord对象
     */
    private AIChatRecord mapResultSetToChatRecord(ResultSet rs) throws SQLException {
        AIChatRecord record = new AIChatRecord();
        record.setId(rs.getString("id"));
        record.setUserId(rs.getString("user_id"));
        record.setUserName(rs.getString("user_name"));
        record.setModelId(rs.getString("model_id"));
        record.setQuestion(rs.getString("question"));
        record.setAnswer(rs.getString("answer"));
        record.setStatus(rs.getInt("status"));
        record.setErrorMessage(rs.getString("error_message"));
        record.setTokensUsed(rs.getInt("tokens_used"));
        record.setCreateTime(rs.getTimestamp("create_time"));
        return record;
    }
    
    /**
     * 将ResultSet映射为AIUsageLimit对象
     */
    private AIUsageLimit mapResultSetToUsageLimit(ResultSet rs) throws SQLException {
        AIUsageLimit limit = new AIUsageLimit();
        limit.setId(rs.getString("id"));
        limit.setUserId(rs.getString("user_id"));
        limit.setDailyLimit(rs.getInt("daily_limit"));
        limit.setUsedToday(rs.getInt("used_today"));
        limit.setLastResetDate(rs.getDate("last_reset_date"));
        limit.setCreateTime(rs.getTimestamp("create_time"));
        limit.setUpdateTime(rs.getTimestamp("update_time"));
        return limit;
    }
}